import phe as paillier
import math
import random

# Generate public and private keys
public_key, private_key = paillier.generate_paillier_keypair()
# Maximum number of bits that can be encrypted
num_of_bits = math.floor(math.log(public_key.n, 2))
# Precision and padding
precision, padding = 31, 1
# Size of a single data item
data_size = precision + padding
# Number of data items that can be packed in a single encryption
num_of_items = math.floor(num_of_bits / data_size)
# Upper and lower bounds for input numbers
LOG2_BOUND = 20
bound = 2 ** LOG2_BOUND


def get_constant() -> tuple[int, int]:
    """Returns the constants used for encryption and decryption."""
    return num_of_items, bound


def get_encrypted_bits(r_list: list[float]) -> paillier.EncryptedNumber:
    """Encrypts a list of floating-point numbers and packs them into a single encryption."""
    pack_bits = ''
    for i, r in enumerate(r_list):
        # Convert float to int
        int_r = math.floor((r + bound) * 2 ** (precision - LOG2_BOUND - 1))
        # Pack data in encryption
        pack_bits += f"{int_r:0{precision}b}"
        if i != num_of_items - 1:
            zero_pad = '0'
            pack_bits += zero_pad.zfill(padding)
    return public_key.encrypt(int('0b' + pack_bits, 2))


def get_decrypted_number(pack_ciphertext: paillier.EncryptedNumber, n: int) -> list[float]:
    """Decrypts a packed ciphertext and returns the list of floating-point numbers."""
    total_sum_bits = bin(private_key.decrypt(pack_ciphertext))[2:].zfill(num_of_items * data_size)
    sum_list = [(int(total_sum_bits[i:i+data_size], 2) / 2 ** (precision - LOG2_BOUND - 1)) - n * bound
                for i in range(0, num_of_items * data_size, data_size)]
    return sum_list


if __name__ == '__main__':
    t, bound = get_constant()
    # Generate a list of random floating-point numbers
    r_list = [random.uniform(-bound, bound) for _ in range(t)]
    print(r_list)
    batch_encrypted1 = get_encrypted_bits(r_list)
    batch_decrypted1 = get_decrypted_number(batch_encrypted1, 1)
    print(batch_decrypted1)